#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "sst/SST.h"

#include "hyper/Hyper.h"

#include <iostream>

namespace mlir {
#define GEN_PASS_DEF_CONVERTCRYPTOTOARM
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir

using namespace mlir;
using namespace vector;
using namespace mlir::arith;

//===----------------------------------------------------------------------===//
// Operation conversion
//===----------------------------------------------------------------------===//

namespace {

class HyperSSTMemcpyLowering : public OpRewritePattern<hyper::MemcpyOp> {
public:
  using OpRewritePattern<hyper::MemcpyOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(hyper::MemcpyOp op,
                                PatternRewriter &rewriter) const override {
    StringRef dir = op.getCopyDirection();
    Value src = op.getSrc();
    Value dst = op.getDst();
    // auto srcOp = src.getDefiningOp();
    // auto dstOp = dst.getDefiningOp();
    auto loc = op.getLoc();

    Value cpy_size =
        rewriter.create<memref::DimOp>(loc, /*source=*/src, /*index=*/0);
    // // get dialect name of src and dst ops
    // StringRef srcDialect = srcOp->getDialect()->getNamespace();
    // StringRef dstDialect = dstOp->getDialect()->getNamespace();

    // // only support data movement between host and device
    // assert(((!srcDialect.equals("hyper") && dstDialect.equals("hyper")) ||
    //        (srcDialect.equals("hyper") && !dstDialect.equals("hyper"))) &&
    //         "only support data movement between host and device");

    // TODO: get the device type from src and dst ops without the need of copy
    // direction
    if (dir.equals("cpu2cpu")) {
      rewriter.replaceOpWithNewOp<memref::CopyOp>(op, /*source=*/src,
                                                  /*target=*/dst);
      //   // erase 'hyper.dealloc' op if its memref operand is to be deleted
      //   for (auto user : dstOp->getUsers()) {
      //     if (auto deallocOp = dyn_cast<hyper::DeallocOp>(user)) {
      //       rewriter.eraseOp(deallocOp);
      //     }
      //   }
      //   // replace 'hyper.alloc' op with its memref operand
      //   rewriter.replaceOp(dstOp, {src});
      //   rewriter.eraseOp(op);
      // } else {
      //   // erase 'hyper.dealloc' op if its memref operand is to be deleted
      //   for (auto user : srcOp->getUsers()) {
      //     if (auto deallocOp = dyn_cast<hyper::DeallocOp>(user)) {
      //       rewriter.eraseOp(deallocOp);
      //     }
      //   }
      //   // replace 'hyper.alloc' op with its memref operand
      //   rewriter.replaceOp(srcOp, {dst});
      //   rewriter.eraseOp(op);
      // }
    } else if (dir.equals("cpu2gpu") || dir.equals("cpu2cgra")) {
      rewriter.replaceOpWithNewOp<sst::MemcpyOp>(
          op, /*dst=*/dst, /*src=*/src, /*size=*/cpy_size,
          /*dir=*/
          rewriter.create<arith::ConstantOp>(
              loc, rewriter.getI1Type(),
              rewriter.getIntegerAttr(rewriter.getI1Type(), APInt(1, 0))));
    } else if (dir.equals("gpu2cpu")) {
      rewriter.replaceOpWithNewOp<sst::MemcpyOp>(
          op, /*dst=*/dst, /*src=*/src, /*size=*/cpy_size,
          /*dir=*/
          rewriter.create<arith::ConstantOp>(
              loc, rewriter.getI1Type(),
              rewriter.getIntegerAttr(rewriter.getI1Type(), APInt(1, 1))));
    } else {
      return rewriter.notifyMatchFailure(
          op, "unknow copy direction for hyper.memcpy");
    }

    return success();
  }
};

} // namespace

void populateLowerHyperSSTMemcpyPatterns(RewritePatternSet &patterns) {
  patterns.add<HyperSSTMemcpyLowering>(patterns.getContext());
}

namespace {
class LowerHyperSSTMemcpyPass
    : public PassWrapper<LowerHyperSSTMemcpyPass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerHyperSSTMemcpyPass)
  LowerHyperSSTMemcpyPass() = default;
  LowerHyperSSTMemcpyPass(const LowerHyperSSTMemcpyPass &) {}

  StringRef getArgument() const final { return "lower-hyper-sstmemcpy"; }
  StringRef getDescription() const final {
    return "Lower hyper.memcpy op to device-specific ops.";
  }

  void runOnOperation() override;

  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<func::FuncDialect, mlir::LLVM::LLVMDialect,
                    memref::MemRefDialect, arith::ArithDialect,
                    vector::VectorDialect, affine::AffineDialect,
                    sst::SSTDialect, scf::SCFDialect>();
  }
};
} // end anonymous namespace.

void LowerHyperSSTMemcpyPass::runOnOperation() {
  MLIRContext *context = &getContext();
  ModuleOp module = getOperation();

  ConversionTarget target(*context);
  target.addLegalDialect<func::FuncDialect, mlir::LLVM::LLVMDialect,
                         memref::MemRefDialect, arith::ArithDialect,
                         affine::AffineDialect, sst::SSTDialect,
                         vector::VectorDialect, scf::SCFDialect>();
  target.addLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>();

  RewritePatternSet patterns(context);
  populateLowerHyperSSTMemcpyPatterns(patterns);

  if (failed(applyPartialConversion(module, target, std::move(patterns))))
    signalPassFailure();
}

void registerLowerHyperSSTMemcpyPass() {
  PassRegistration<LowerHyperSSTMemcpyPass>();
}

std::unique_ptr<Pass> createLowerHyperSSTMemcpyPass() {
  return std::make_unique<LowerHyperSSTMemcpyPass>();
}
